# -*- coding: utf-8 -*-
# @Time    : 2023/9/20 17:17
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : idealLGB.py
# @Software: PyCharm
import typing

import numpy
import pandas


class IdealLGB:
    BIGNUM = 10086

    class Colname:
        dz, B, z = 'dz, B, z'.split(', ')

    def __init__(self, z_magnet_start, dzs: typing.List[float], Bs: typing.List[float]):
        self.z0 = z_magnet_start
        self.data = pandas.DataFrame({
            IdealLGB.Colname.dz: [self.BIGNUM, *dzs, self.BIGNUM, self.BIGNUM],  # 为了涵盖粒子可能经过的0磁场位置
            IdealLGB.Colname.B: [0, *Bs, 0, 0]
        })
        self.data[IdealLGB.Colname.z] = [self.data[IdealLGB.Colname.dz][:i].sum() + self.z0 - self.BIGNUM for i in
                                         range(self.data.shape[0])]  # 每个台阶的起始位置


class Traj:
    class Colname:
        x, z = 'x, z'.split(', ')

    def __init__(self, xs, zs):
        self.data = pandas.DataFrame({'x': xs, 'z': zs})

    def get_ds(self):
        """
        获取相邻轨迹点之间的距离
        :return:
        """
        ds = (numpy.diff(self.data, axis=0) ** 2).sum(axis=1) ** 0.5
        return numpy.append(ds, 0)

    def get_ideal_B(self, idealLGB: IdealLGB):
        """
        获取沿粒子轨迹上每个点的理想磁场大小
        :param idealLGB:
        :return:
        """
        cut = pandas.cut(self.data.z, idealLGB.data[IdealLGB.Colname.z], labels=False).astype(int)
        B = numpy.array([idealLGB.data[IdealLGB.Colname.B][cut[i]] for i in range(len(cut))])
        return B, cut

    def get_integral(self, B_along_traj, ds):
        return (B_along_traj * ds).sum()


if __name__ == '__main__':
    ideal_traj_df = pandas.read_csv(r"F:\GPTfiles\LGB\integration\pys\temp.ideal.csv")
    optimized_traj_df = pandas.read_csv(r"F:\GPTfiles\LGB\integration\pys\temp.optimized.csv",header=None)
    idealLGB = IdealLGB(0, [0.229869186700053, 0.229941263822213,
                            0.229990187110817, 0.229977996014191,
                            0.229957875298085],
                        [-0.333400762, -0.385611315, -0.538732484,
                         -0.792828097, -0.427035091, ])
    ideal_traj = Traj(ideal_traj_df['x(mm)'] * 1e-3, ideal_traj_df['z(mm)'] * 1e-3, )
    ideal_intB = ideal_traj.get_integral(ideal_traj.get_ideal_B(idealLGB)[0], ideal_traj.get_ds())
    optimized_traj = Traj(optimized_traj_df[1] * 1e-3, optimized_traj_df[0] * 1e-3, )
    optimized_intB = optimized_traj.get_integral(optimized_traj.get_ideal_B(idealLGB)[0], optimized_traj.get_ds())
